﻿using Microsoft.ML.OnnxRuntime;
using Microsoft.ML.OnnxRuntime.Tensors;
using OpenCvSharp;
using System;
using System.Collections.Generic;
using System.Linq;
using System.Text;
using System.Threading.Tasks;

namespace Common.plate.common
{
    /// <summary>
    /// OCR 文字识别器(识别车牌号码)
    /// </summary>
    public class TextRecognizer : IDisposable
    {
        private InferenceSession _session; // ONNX推理会话
        private int _inputWidth; // 模型输入宽度
        private int _inputHeight; // 模型输入高度
        private List<char> _characters; // 字符集（如"0123456789ABCDEFG...京沪粤..."）


        /// <summary>
        /// 初始化OCR识别器
        /// </summary>
        /// <param name="modelPath">OCR模型路径</param>
        /// <param name="charsPath">字符集文件路径（每行一个字符）</param>
        public TextRecognizer(string modelPath, string charsPath)
        {
            // 加载ONNX模型
            var options = new SessionOptions()
            {
                LogSeverityLevel = OrtLoggingLevel.ORT_LOGGING_LEVEL_WARNING // 只输出警告日志，减少冗余
            };
            _session = new InferenceSession(modelPath, options);

            // 获取模型输入尺寸（假设输入形状为[1, 3, H, W]）
            var inputMeta = _session.InputMetadata.Values.First();
            int[]? shape = inputMeta.Dimensions;

            _inputHeight = shape[2];//(640)
            _inputWidth = shape[3];//(640)

            // 加载字符集（用于将模型输出的索引转换为文字）
            // 加载字符集（用于将模型输出的索引转换为文字）
            _characters = System.IO.File.ReadAllLines(charsPath)
                .Select(line => line.Trim()[0]) // 每行取第一个字符
                .ToList();
            _characters.Insert(0, ' '); // 插入空白字符（CTC解码用）
        }

        /// <summary>
        /// 识别图像中的文字（车牌）
        /// </summary>
        /// <param name="images">车牌图像列表（裁剪后的ROI）</param>
        /// <returns>识别结果列表</returns>
        public List<OcrResult> Recognize(List<Mat> images)
        {
            var results = new List<OcrResult>();
            foreach (var img in images)
            {
                // 获取模型输入节点名称
                var inputName = _session.InputNames[0];
                // 1. 图像预处理
                var inputTensor = Preprocess(img);

                // 2. 模型推理
                var inputs = new List<NamedOnnxValue>
                {
                NamedOnnxValue.CreateFromTensor(inputName, inputTensor) // 输入名称需与模型一致
                };
                using var outputs = _session.Run(inputs);
                var outputTensor = outputs.First().AsTensor<float>(); // 形状：[1, seq_len, num_classes]

                // 3. 解码结果（CTC解码）
                var (text, score) = Decode(outputTensor);
                results.Add(new OcrResult { Text = text, Score = score });
            }
            return results;
        }

        /// <summary>
        /// 图像预处理：缩放、归一化、转换为张量
        /// </summary>
        private DenseTensor<float> Preprocess(Mat img)
        {
            // 保存原始尺寸
            int originalWidth = img.Cols;
            int originalHeight = img.Rows;

            // 计算缩放比例（保持宽高比）
            float scale = Math.Min((float)_inputWidth / originalWidth, (float)_inputHeight / originalHeight);
            int newWidth = (int)(originalWidth * scale);
            int newHeight = (int)(originalHeight * scale);
            // 重新计算以确保尺寸准确
            newWidth = Math.Min(newWidth, _inputWidth);
            newHeight = Math.Min(newHeight, _inputHeight);

            // 缩放图像
            Mat resized = new Mat();
            Cv2.Resize(img, resized, new Size(newWidth, newHeight));

            // 计算填充（使图像尺寸符合模型输入，避免拉伸）
            int padW = (_inputWidth - newWidth) / 2;
            int padH = (_inputHeight - newHeight) / 2;    

            // 填充（上下左右填充黑色）
            Mat padded = new Mat();
            Cv2.CopyMakeBorder(resized, padded, padH, _inputHeight - newHeight - padH,
                              padW, _inputWidth - newWidth - padW, BorderTypes.Constant, Scalar.Black);

            // 转换为灰度图（如果模型是单通道输入）
            if (_inputChannels == 1)
            {
                Cv2.CvtColor(padded, resized, ColorConversionCodes.BGR2GRAY);
            }
            else
            {
                // 转换为RGB（如果模型是3通道输入）
                Cv2.CvtColor(padded, resized, ColorConversionCodes.BGR2RGB);
            }

            // 归一化（除以255，转换为float）
            resized.ConvertTo(resized, MatType.CV_32FC(_inputChannels), 1.0 / 255.0);

            // 转换为张量（形状：[1, C, H, W]）
            int channels = _inputChannels;
            var tensor = new DenseTensor<float>(new[] { 1, channels, _inputHeight, _inputWidth });
            for (int c = 0; c < channels; c++)
            {
                for (int y = 0; y < _inputHeight; y++)
                {
                    for (int x = 0; x < _inputWidth; x++)
                    {
                        tensor[0, c, y, x] = channels == 1
                            ? (float)resized.At<float>(y, x)
                            : (float)resized.At<Vec3f>(y, x)[c];
                    }
                }
            }
            return tensor;
        }

        /// <summary>
        /// CTC解码：将模型输出的概率分布转换为文字
        /// </summary>
        private (string text, float score) Decode(Tensor<float> output)
        {
            var shape = output.Dimensions.ToArray();
            int seqLen = output.Dimensions[1]; // 序列长度,字符位置数（如7位车牌对应7个时间步）
            int numClasses = output.Dimensions[2]; // 字符类别数

            // 1. 取每个时间步的最大概率索引
            var indices = new List<int>();
            var scores = new List<float>();
            for (int t = 0; t < seqLen; t++)
            {
                float maxScore = 0;
                int maxIndex = 0;
                for (int c = 0; c < numClasses; c++)
                {
                    float score = output[0, t, c];
                    if (score > maxScore)
                    {
                        maxScore = score;
                        maxIndex = c;
                    }
                }
                indices.Add(maxIndex);
                scores.Add(maxScore);
            }

            // 2. 去重（CTC解码逻辑：连续相同索引只保留一个，跳过空白字符）
            var decodedIndices = new List<int>();//有效的索引
            int prev = -1;
            foreach (int idx in indices)
            {
                if (idx != prev && idx != 0) // 0是空白字符索引
                {
                    if(idx < _characters.Count)
                    {
                        decodedIndices.Add(idx);
                        prev = idx;
                    }                    
                }
            }

            // 3. 转换为文字并计算平均置信度
            string text = new string(decodedIndices.Select(idx => _characters[idx]).ToArray());
            float avgScore = decodedIndices.Count == 0
                ? 0
                : scores.Where((_, i) => indices[i] != 0 && indices[i] == decodedIndices[0]).Average();

            return (text, avgScore);
        }

        public void Dispose()
        {
            _session?.Dispose();
        }

        // 模型输入通道数（根据实际模型调整，1=灰度，3=彩色）
        private int _inputChannels => _session.InputMetadata.Values.First().Dimensions[1] == 1 ? 1 : 3;
        /// <summary>
        /// OCR识别结果
        /// </summary>
        public class OcrResult
        {
            public string Text { get; set; } // 识别到的文字
            public float Score { get; set; } // 平均置信度
        }
    }
}
